package net.sourceforge.jsocks;

import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.UnknownHostException;

/**
 * SOCKS5 request/response message.
 */

public class Socks5Message extends ProxyMessage {
    // SOCKS5 constants
    public static final int SOCKS_VERSION = 5;
    public static final int SOCKS_ATYP_IPV4 = 0x1; // Where is 2??
    public static final int SOCKS_ATYP_DOMAINNAME = 0x3; // !!!!rfc1928
    public static final int SOCKS_ATYP_IPV6 = 0x4;
    public static final int SOCKS_IPV6_LENGTH = 16;
    static boolean doResolveIP = true;
    /**
     * Address type of given message
     */
    public int addrType;
    byte[] data;

    /**
     * Initialises Message from the stream. Reads server response from given
     * stream.
     *
     * @param in Input stream to read response from.
     * @throws SocksException If server response code is not SOCKS_SUCCESS(0), or if any
     *                        error with protocol occurs.
     * @throws IOException    If any error happens with I/O.
     */
    public Socks5Message(InputStream in) throws SocksException, IOException {
        this(in, true);
    }

    /**
     * Initialises Message from the stream. Reads server response or client
     * request from given stream.
     *
     * @param in         Input stream to read response from.
     * @param clinetMode If true read server response, else read client request.
     * @throws SocksException If server response code is not SOCKS_SUCCESS(0) and reading
     *                        in client mode, or if any error with protocol occurs.
     * @throws IOException    If any error happens with I/O.
     */
    public Socks5Message(InputStream in, boolean clientMode)
            throws SocksException, IOException {
        read(in, clientMode);
    }

    /**
     * Server error response.
     *
     * @param cmd Error code.
     */
    public Socks5Message(int cmd) {
        super(cmd, null, 0);
        data = new byte[3];
        data[0] = SOCKS_VERSION; // Version.
        data[1] = (byte) cmd; // Reply code for some kind of failure.
        data[2] = 0; // Reserved byte.
    }

    /**
     * Construct client request or server response.
     *
     * @param cmd - Request/Response code.
     * @param ip  - IP field.
     * @paarm port - port field.
     */
    public Socks5Message(int cmd, InetAddress ip, int port) {
        super(cmd, ip, port);
        this.host = ip == null ? "0.0.0.0" : ip.getHostName();
        this.version = SOCKS_VERSION;

        byte[] addr;

        if (ip == null) {
            addr = new byte[4];
            addr[0] = addr[1] = addr[2] = addr[3] = 0;
        } else
            addr = ip.getAddress();

        addrType = addr.length == 4 ? SOCKS_ATYP_IPV4 : SOCKS_ATYP_IPV6;

        data = new byte[6 + addr.length];
        data[0] = (byte) SOCKS_VERSION; // Version
        data[1] = (byte) command; // Command
        data[2] = (byte) 0; // Reserved byte
        data[3] = (byte) addrType; // Address type

        // Put Address
        System.arraycopy(addr, 0, data, 4, addr.length);
        // Put port
        data[data.length - 2] = (byte) (port >> 8);
        data[data.length - 1] = (byte) (port);
    }

    /**
     * Construct client request or server response.
     *
     * @param cmd      - Request/Response code.
     * @param hostName - IP field as hostName, uses ADDR_TYPE of HOSTNAME.
     * @paarm port - port field.
     */
    public Socks5Message(int cmd, String hostName, int port) {
        super(cmd, null, port);
        this.host = hostName;
        this.version = SOCKS_VERSION;

        // System.out.println("Doing ATYP_DOMAINNAME");

        addrType = SOCKS_ATYP_DOMAINNAME;
        byte addr[] = hostName.getBytes();

        data = new byte[7 + addr.length];
        data[0] = (byte) SOCKS_VERSION; // Version
        data[1] = (byte) command; // Command
        data[2] = (byte) 0; // Reserved byte
        data[3] = (byte) SOCKS_ATYP_DOMAINNAME; // Address type
        data[4] = (byte) addr.length; // Length of the address

        // Put Address
        System.arraycopy(addr, 0, data, 5, addr.length);
        // Put port
        data[data.length - 2] = (byte) (port >> 8);
        data[data.length - 1] = (byte) (port);
    }

    /**
     * Wether to resolve hostIP returned from SOCKS server that is wether to
     * create InetAddress object from the hostName string
     */
    static public boolean resolveIP() {
        return doResolveIP;
    }

    /*
     * private static final void debug(String s){ if(DEBUG) System.out.print(s);
     * } private static final boolean DEBUG = false;
     */

    /**
     * Wether to resolve hostIP returned from SOCKS server that is wether to
     * create InetAddress object from the hostName string
     *
     * @param doResolve Wether to resolve hostIP from SOCKS server.
     * @return Previous value.
     */
    static public boolean resolveIP(boolean doResolve) {
        boolean old = doResolveIP;
        doResolveIP = doResolve;
        return old;
    }

    /**
     * Returns IP field of the message as IP, if the message was created with
     * ATYP of HOSTNAME, it will attempt to resolve the hostname, which might
     * fail.
     *
     * @throws UnknownHostException if host can't be resolved.
     */
    @Override
    public InetAddress getInetAddress() throws UnknownHostException {
        if (ip != null)
            return ip;

        return (ip = InetAddress.getByName(host));
    }

    /**
     * Initialises Message from the stream. Reads server response from given
     * stream.
     *
     * @param in Input stream to read response from.
     * @throws SocksException If server response code is not SOCKS_SUCCESS(0), or if any
     *                        error with protocol occurs.
     * @throws IOException    If any error happens with I/O.
     */
    @Override
    public void read(InputStream in) throws SocksException, IOException {
        read(in, true);
    }

    /**
     * Initialises Message from the stream. Reads server response or client
     * request from given stream.
     *
     * @param in         Input stream to read response from.
     * @param clinetMode If true read server response, else read client request.
     * @throws SocksException If server response code is not SOCKS_SUCCESS(0) and reading
     *                        in client mode, or if any error with protocol occurs.
     * @throws IOException    If any error happens with I/O.
     */
    @Override
    public void read(InputStream in, boolean clientMode) throws SocksException,
            IOException {
        data = null;
        ip = null;

        DataInputStream di = new DataInputStream(in);

        version = di.readUnsignedByte();
        command = di.readUnsignedByte();
        if (clientMode && command != 0)
            throw new SocksException(command);

        @SuppressWarnings("unused")
        int reserved = di.readUnsignedByte();
        addrType = di.readUnsignedByte();

        byte addr[];

        switch (addrType) {
            case SOCKS_ATYP_IPV4:
                addr = new byte[4];
                di.readFully(addr);
                host = bytes2IPV4(addr, 0);
                break;
            case SOCKS_ATYP_IPV6:
                addr = new byte[SOCKS_IPV6_LENGTH];// I believe it is 16 bytes,huge!
                di.readFully(addr);
                host = bytes2IPV6(addr, 0);
                break;
            case SOCKS_ATYP_DOMAINNAME:
                // System.out.println("Reading ATYP_DOMAINNAME");
                addr = new byte[di.readUnsignedByte()];// Next byte shows the length
                di.readFully(addr);
                host = new String(addr);
                break;
            default:
                throw (new SocksException(Proxy.SOCKS_JUST_ERROR));
        }

        port = di.readUnsignedShort();

        if (addrType != SOCKS_ATYP_DOMAINNAME && doResolveIP) {
            try {
                ip = InetAddress.getByName(host);
            } catch (UnknownHostException uh_ex) {
            }
        }
    }

    /**
     * Returns string representation of the message.
     */
    @Override
    public String toString() {
        String s = "Socks5Message:" + "\n" + "VN   " + version + "\n" + "CMD  "
                + command + "\n" + "ATYP " + addrType + "\n" + "ADDR " + host
                + "\n" + "PORT " + port + "\n";
        return s;
    }

    /**
     * Writes the message to the stream.
     *
     * @param out Output stream to which message should be written.
     */
    @Override
    public void write(OutputStream out) throws SocksException, IOException {
        if (data == null) {
            Socks5Message msg;

            if (addrType == SOCKS_ATYP_DOMAINNAME)
                msg = new Socks5Message(command, host, port);
            else {
                if (ip == null) {
                    try {
                        ip = InetAddress.getByName(host);
                    } catch (UnknownHostException uh_ex) {
                        throw new SocksException(Proxy.SOCKS_JUST_ERROR);
                    }
                }
                msg = new Socks5Message(command, ip, port);
            }
            data = msg.data;
        }
        out.write(data);
    }

}
